/*
 * @brief: the implementation of matrix multiplication using FMAC
 */
#include "benchmark_config.h"
#include "gpu_timer.h"
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#define N 1024 // Matrix size (N x N)
#define BLOCK_SIZE 32
#define NUM_RUNS 1000 // Number of times to run the kernel for averaging

// Define v16F32 as a vector of 16 bfloat16 values
typedef struct {
  float data[16];
} v16F32;

// Kernel for matrix multiplication using v16F32
__global__ void matrixMultiplyKernel(const v16F32 *A, const v16F32 *B,
                                     v16F32 *C, int n) {
  int row = blockIdx.y * blockDim.y + threadIdx.y;
  int col = blockIdx.x * blockDim.x + threadIdx.x;

  if (row < n && col < n / 16) {
    v16F32 sum;
    for (int i = 0; i < 16; ++i) {
      sum.data[i] = 0.0f;
    }

    for (int i = 0; i < n / 16; ++i) {
      v16F32 a = A[row * (n / 16) + i];
      v16F32 b = B[i * (n / 16) + col];

      // Perform element-wise multiplication and addition
      for (int j = 0; j < 16; ++j) {
        for (int k = 0; k < 16; ++k) {
          sum.data[j] = (sum.data[j] + (a.data[k] * b.data[(j + k) % 16]));
        }
      }
    }
    C[row * (n / 16) + col] = sum;
  }
}

// Function to initialize a matrix with random v16F32 values
void initMatrix(v16F32 *matrix, int size) {
  for (int i = 0; i < size; i++) {
    for (int j = 0; j < 16; j++) {
      float randVal = static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
      matrix[i].data[j] = randVal;
    }
  }
}

// Function to verify the result
bool verifyResult(const v16F32 *A, const v16F32 *B, const v16F32 *C, int n) {
  const float epsilon = 1e-2f; // Larger epsilon due to F32 precision
  for (int i = 0; i < n; ++i) {
    for (int j = 0; j < n / 16; ++j) {
      v16F32 sum;
      for (int k = 0; k < 16; ++k) {
        sum.data[k] = float(0.0f);
      }
      for (int k = 0; k < n / 16; ++k) {
        v16F32 a = A[i * (n / 16) + k];
        v16F32 b = B[k * (n / 16) + j];
        for (int l = 0; l < 16; ++l) {
          for (int m = 0; m < 16; ++m) {
            sum.data[l] = (sum.data[l] + (a.data[m] * b.data[(l + m) % 16]));
          }
        }
      }
      v16F32 c = C[i * (n / 16) + j];
      for (int k = 0; k < 16; ++k) {
        if (fabs(float(sum.data[k]) - float(c.data[k])) > epsilon) {
          return false;
        }
      }
    }
  }
  return true;
}

int main() {
  int size = N * N / 16;
  size_t bytes = size * sizeof(v16F32);

  // Allocate host memory
  v16F32 *h_A = (v16F32 *)malloc(bytes);
  v16F32 *h_B = (v16F32 *)malloc(bytes);
  v16F32 *h_C = (v16F32 *)malloc(bytes);

  // Initialize host matrices
  initMatrix(h_A, size);
  initMatrix(h_B, size);

  // Allocate device memory
  v16F32 *d_A, *d_B, *d_C;
  cudaMalloc(&d_A, bytes);
  cudaMalloc(&d_B, bytes);
  cudaMalloc(&d_C, bytes);

  // Copy host memory to device
  cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
  cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice);

  // Setup execution configuration
  dim3 threadsPerBlock(BLOCK_SIZE, BLOCK_SIZE);
  dim3 blocksPerGrid((N / 16 + BLOCK_SIZE - 1) / BLOCK_SIZE,
                     (N + BLOCK_SIZE - 1) / BLOCK_SIZE);

  // warm up
  int warmup = WARMUPNUM;
  for (int i = 0; i < warmup; i++)
    matrixMultiplyKernel<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);
  cudaDeviceSynchronize();

  float totalTime = 0.0f;
  GpuTimer timer;

  // Run kernel multiple times and calculate average execution time
  for (int i = 0; i < NUM_RUNS; i++) {
    timer.start();
    matrixMultiplyKernel<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);
    timer.stop_and_wait();
    totalTime += timer.duration(1);
  }

  float averageTime = totalTime / NUM_RUNS;
  printf("Average kernel execution time: %f ms\n", averageTime);

  double ips = N * N * N * 2. / 1000. / (double)averageTime;
  printf("Average matrixMultiplyKernel flops: %lf\n", ips);

  // Copy result back to host
  cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);

  // Verify the result
  bool correct = verifyResult(h_A, h_B, h_C, N);
  printf("Matrix multiplication %s\n", correct ? "PASS" : "FAIL");

  // Clean up
  cudaFree(d_A);
  cudaFree(d_B);
  cudaFree(d_C);
  free(h_A);
  free(h_B);
  free(h_C);

  if (!correct) {
    return -1;
  }

  return 0;
}
